from src.common.config import config
from src.common.logger import getLogger
from langchain_core.documents import Document
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient, models
from langchain_text_splitters import RecursiveCharacterTextSplitter

logger = getLogger()

class VectorStore:

    def __init__(self):
        config_dict = config.parse_config_key(["qdrant"])
        logger.info(f"VectorStore __init__ config_dict: {config_dict}")
        self.host = config_dict.get("host", "")
        self.port = config_dict.get("port", 0)
        self.timeout = config_dict.get("timeout", 0)
        self.vector_size = config_dict.get("vector_size", 0)

    def new_vector_client(self):
        qdrant_client = QdrantClient(
            host = self.host,
            port = self.port,
            timeout = self.timeout
        )
        logger.info(f"VectorStore new_vector_client qdrant_client: {qdrant_client}")
        return qdrant_client

    def new_vector_store(self, embedding, collection_name):
        logger.info(f"VectorStore new_vector_store collection_name: {collection_name}")
        qdrant_client = self.new_vector_client()
        exist = qdrant_client.collection_exists(collection_name)
        logger.info(f"VectorStore new_vector_store collection_name exist: {exist}")
        if not exist:
            qdrant_client.create_collection(
                collection_name = collection_name,
                vectors_config = models.VectorParams(size = self.vector_size, distance = models.Distance.COSINE)
            )
        logger.info(f"VectorStore new_vector_store collection_name info: {qdrant_client.info()}")

        vector_store = QdrantVectorStore(client = qdrant_client, collection_name = collection_name, embedding = embedding)
        return vector_store

    def add_document_vector(self, embedding, collection_name, content):
        logger.info(f"VectorStore add_document_vector collection_name: {collection_name}")
        vector_store = self.new_vector_store(embedding, collection_name)
        logger.info(f"VectorStore add_document_vector vector_store {collection_name} finished")

        text_splitter = RecursiveCharacterTextSplitter(chunk_size = 768, chunk_overlap = 50)
        split_texts = text_splitter.split_text(content)
        documents = [Document(page_content = text, metadata = { "doc_id": i }) for i, text in enumerate(split_texts)]
        logger.info(f"VectorStore add_document_vector documents len: {len(documents)}")

        vector_store.add_documents(documents)
        logger.info(f"VectorStore add_document_vector vector_store {collection_name} finished")
